import numpy as np
import matplotlib.pyplot as plt


# Define the Schaffer Function N. 4
def schaffer_n4(x, y):
    return 0.5 + (np.cos(np.sin(abs(x ** 2 - y ** 2))) ** 2 - 0.5) / (1 + 0.001 * (x ** 2 + y ** 2)) ** 2


# Whale Optimization Algorithm (WOA)
class WhaleOptimizationAlgorithm:
    def __init__(self, population_size, dim, max_iter):
        self.population_size = population_size
        self.dim = dim
        self.max_iter = max_iter
        self.pos = np.random.uniform(-10, 10, (population_size, dim))
        self.best_pos = None
        self.best_score = float('inf')
        self.scores = np.zeros(population_size)
        self.history = []

    def optimize(self):
        for i in range(self.max_iter):
            for j in range(self.population_size):
                self.scores[j] = schaffer_n4(self.pos[j, 0], self.pos[j, 1])

                if self.scores[j] < self.best_score:
                    self.best_score = self.scores[j]
                    self.best_pos = self.pos[j].copy()

            a = 2 - i * (2 / self.max_iter)  # a decreases linearly from 2 to 0

            for j in range(self.population_size):
                r = np.random.rand()  # random number in [0, 1)
                A = 2 * a * r - a
                C = 2 * r
                b = 1
                l = (np.random.rand() - 0.5) * 2
                p = np.random.rand()

                if p < 0.5:
                    if abs(A) < 1:
                        D = abs(C * self.best_pos - self.pos[j])
                        self.pos[j] = self.best_pos - A * D
                    else:
                        random_agent = np.random.randint(0, self.population_size)
                        D = abs(C * self.pos[random_agent] - self.pos[j])
                        self.pos[j] = self.pos[random_agent] - A * D
                else:
                    D = abs(self.best_pos - self.pos[j])
                    self.pos[j] = D * np.exp(b * l) * np.cos(2 * np.pi * l) + self.best_pos

            self.history.append(self.best_score)

    def plot_optimization(self):
        plt.plot(self.history)
        plt.title('Whale Optimization Algorithm')
        plt.xlabel('Iteration')
        plt.ylabel('Best Score')
        plt.show()


# Initialize and run the WOA
woa = WhaleOptimizationAlgorithm(population_size=30, dim=2, max_iter=1000)
woa.optimize()
woa.plot_optimization()
